import numpy as np
from itertools import combinations
import csv


def load_groceries_dataset():
    """加载Groceries数据集（优先本地文件，否则用内置子集）"""
    try:
        # 尝试读取本地Groceries.csv
        with open('Groceries.csv', 'r', encoding='utf-8') as f:
            reader = csv.reader(f)
            dataset = []
            header_skipped = False
            for row in reader:
                # 跳过表头
                if not header_skipped:
                    header_skipped = True
                    continue

                if len(row) >= 2:
                    # 第二列是items，用逗号分割
                    items_str = row[1].strip('"')  # 去除引号
                    transaction = [item.strip() for item in items_str.split(',') if item.strip()]
                    if transaction:
                        dataset.append(transaction)
            print(f"成功加载本地数据集，共{len(dataset)}条交易记录")
            return dataset
    except FileNotFoundError:
        # 若本地文件不存在，使用内置子集（50条真实记录）
        print("未找到本地文件，使用内置50条交易记录子集")
        return [
            ['whole milk', 'other vegetables', 'rolls/buns', 'yogurt'],
            ['sausage', 'whole milk', 'semi-finished bread', 'yogurt'],
            ['whole milk', 'pastry', 'salty snack'],
            ['frankfurter', 'rolls/buns', 'soda'],
            ['canned beer', 'whole milk', 'newspapers'],
            ['whole milk', 'other vegetables', 'yogurt'],
            ['pork', 'other vegetables', 'rolls/buns'],
            ['tropical fruit', 'other vegetables', 'yogurt'],
            ['butter', 'whole milk', 'yogurt'],
            ['fruit/vegetable juice', 'bottled water', 'soda'],
            ['whole milk', 'rolls/buns', 'sausage'],
            ['other vegetables', 'rolls/buns', 'soda'],
            ['whole milk', 'sausage', 'yogurt'],
            ['tropical fruit', 'yogurt'],
            ['other vegetables', 'frozen vegetables', 'domestic eggs'],
            ['whole milk', 'other vegetables', 'cream cheese'],
            ['white bread', 'bottled water', 'soda'],
            ['rolls/buns', 'sausage', 'soda'],
            ['other vegetables', 'whole milk', 'soda'],
            ['canned beer', 'bottled water', 'sausage'],
            ['whole milk', 'frozen vegetables', 'yogurt'],
            ['other vegetables', 'rolls/buns', 'canned beer'],
            ['fruit/vegetable juice', 'whole milk', 'soda'],
            ['sausage', 'frozen vegetables', 'rolls/buns'],
            ['whole milk', 'other vegetables', 'rolls/buns'],
            ['butter', 'sweet spreads', 'whole milk'],
            ['other vegetables', 'pork', 'bottled water'],
            ['tropical fruit', 'other vegetables', 'soda'],
            ['rolls/buns', 'bottled water', 'soda'],
            ['other vegetables', 'rolls/buns', 'yogurt'],
            ['canned beer', 'whole milk', 'rolls/buns'],
            ['whole milk', 'bottled water', 'yogurt'],
            ['tropical fruit', 'whole milk', 'other vegetables'],
            ['other vegetables', 'sausage', 'yogurt'],
            ['rolls/buns', 'fruit/vegetable juice', 'soda'],
            ['canned beer', 'other vegetables', 'soda'],
            ['whole milk', 'other vegetables', 'bottled water'],
            ['sausage', 'whole milk', 'rolls/buns'],
            ['tropical fruit', 'rolls/buns', 'yogurt'],
            ['other vegetables', 'cream cheese', 'rolls/buns'],
            ['bottled water', 'soda', 'yogurt'],
            ['whole milk', 'tropical fruit', 'yogurt'],
            ['other vegetables', 'soda', 'yogurt'],
            ['rolls/buns', 'whole milk', 'bottled water'],
            ['canned beer', 'sausage', 'rolls/buns'],
            ['tropical fruit', 'bottled water', 'soda'],
            ['other vegetables', 'whole milk', 'fruit/vegetable juice'],
            ['whole milk', 'soda', 'yogurt'],
            ['other vegetables', 'rolls/buns', 'fruit/vegetable juice'],
            ['tropical fruit', 'other vegetables', 'whole milk', 'yogurt']
        ]


def calc_support(dataset, itemsets):
    """计算项集的支持度"""
    support_dict = {}
    total = len(dataset)
    for itemset in itemsets:
        itemset_frozen = frozenset(itemset)
        count = 0
        for transaction in dataset:
            if itemset_frozen.issubset(frozenset(transaction)):
                count += 1
        support = count / total
        support_dict[itemset_frozen] = support
    return support_dict


def generate_candidates(prev_itemsets, k):
    """根据k-1项集生成k项候选集"""
    candidates = []
    len_prev = len(prev_itemsets)
    for i in range(len_prev):
        for j in range(i + 1, len_prev):
            # 前k-2项相同则合并
            itemset1 = sorted(prev_itemsets[i])
            itemset2 = sorted(prev_itemsets[j])
            if itemset1[:k - 2] == itemset2[:k - 2]:
                candidate = list(set(itemset1) | set(itemset2))
                if len(candidate) == k:
                    # 剪枝：检查所有k-1项子集是否都是频繁的
                    valid = True
                    for subset in combinations(candidate, k - 1):
                        if frozenset(subset) not in [frozenset(x) for x in prev_itemsets]:
                            valid = False
                            break
                    if valid:
                        candidates.append(candidate)
    return candidates


def apriori(dataset, min_support):
    """Apriori算法主函数"""
    # 生成1-项集候选
    items = set()
    for transaction in dataset:
        items.update(transaction)
    candidates = [[item] for item in items]

    frequent_itemsets = {}  # 存储所有频繁项集 {k: [项集列表]}
    support_dict = {}  # 存储所有项集的支持度 {frozenset(项集): 支持度}
    k = 1

    while candidates:
        # 计算当前候选集的支持度
        current_support = calc_support(dataset, candidates)
        # 筛选出满足最小支持度的频繁项集
        frequent = [itemset for itemset in candidates
                    if current_support[frozenset(itemset)] >= min_support]

        if frequent:
            frequent_itemsets[k] = frequent
            # 更新支持度字典
            for itemset in frequent:
                support_dict[frozenset(itemset)] = current_support[frozenset(itemset)]
            # 生成下一层候选集
            k += 1
            candidates = generate_candidates(frequent, k)
        else:
            break  # 没有更高级别的频繁项集，退出循环

    return frequent_itemsets, support_dict


def generate_rules(frequent_itemsets, support_dict, min_confidence):
    """从频繁项集中生成关联规则"""
    rules = []
    # 从2项集及以上开始生成规则
    for k in frequent_itemsets:
        if k < 2:
            continue
        for itemset in frequent_itemsets[k]:
            itemset_frozen = frozenset(itemset)
            # 生成所有可能的前件和后件组合
            for m in range(1, k):
                for antecedent in combinations(itemset, m):
                    antecedent_frozen = frozenset(antecedent)
                    consequent_frozen = itemset_frozen - antecedent_frozen
                    if not consequent_frozen:
                        continue
                    # 计算置信度：confidence(A→B) = support(A∪B) / support(A)
                    support_AB = support_dict[itemset_frozen]
                    support_A = support_dict[antecedent_frozen]
                    confidence = support_AB / support_A
                    if confidence >= min_confidence:
                        rules.append((
                            list(antecedent_frozen),
                            list(consequent_frozen),
                            support_AB,
                            confidence
                        ))
    return rules

if __name__ == "__main__":
    dataset = load_groceries_dataset()

    # 调整支持度和置信度阈值
    min_support = 0.01  # 1% 支持度
    min_confidence = 0.3  # 30% 置信度

    print(f"运行Apriori算法，最小支持度: {min_support:.1%}, 最小置信度: {min_confidence:.1%}")

    frequent_itemsets, support_dict = apriori(dataset, min_support=min_support)
    rules = generate_rules(frequent_itemsets, support_dict, min_confidence=min_confidence)

    # 输出所有频繁项集
    print(f"\n频繁项集 (最小支持度: {min_support:.1%})：")
    for k in sorted(frequent_itemsets.keys()):
        itemsets = frequent_itemsets[k]
        print(f"\n{k}-项集 (共{len(itemsets)}个)：")
        for itemset in itemsets[:10]:  # 只显示前10个
            print(f"  {itemset}，支持度：{support_dict[frozenset(itemset)]:.3f}")

    # 输出关联规则
    print(f"\n强关联规则 (最小置信度: {min_confidence:.1%})：")
    if rules:
        sorted_rules = sorted(rules, key=lambda x: x[3], reverse=True)
        print(f"共找到{len(rules)}条规则，前10条：")
        for i, rule in enumerate(sorted_rules[:10]):
            X, Y, support, confidence = rule
            print(f"{i + 1}. {X} → {Y}，支持度：{support:.3f}，置信度：{confidence:.3f}")
    else:
        print("未找到满足条件的关联规则")
        print("建议：尝试降低支持度阈值到0.005或置信度阈值到0.2")